package com.shujia.spark.mllib

import org.apache.spark.internal.Logging
import org.apache.spark.ml.classification.{LogisticRegression, LogisticRegressionModel}
import org.apache.spark.sql.{DataFrame, Dataset, Row, SparkSession}
import org.slf4j.Logger

object Demo02PersonIndexTrain extends Logging {

  val logger: Logger = log

  def main(args: Array[String]): Unit = {

    /**
     * spark-submit --master local[*] --conf spark.sql.shuffle.partitions=2 --class com.shujia.spa
     * rk.mllib.Demo02PersonIndexTrain spark-1.0.jar
     */
    if (args.length != 2) {
      logger.error("请检查参数的个数，需要两个参数：数据输入路径，模型保存路径")
      return
    }

    val inputPath: String = args(0)
    val outputPath: String = args(1)
    logger.info(s"当前的输入路径为:$inputPath")
    logger.info(s"当前的模型保存路径为:$outputPath")

    // 1、进行数据特征工程，将数据转换成SparkMLLib能够识别的向量
    val spark: SparkSession = SparkSession
      .builder()
      .appName(this.getClass.getSimpleName.replace("$", ""))
      //      .master("local") // 设置运行的方式
      .getOrCreate()

    // SparkSQL加载libsvm格式的数据
    val libsvmDF: DataFrame = spark
      .read
      .format("libsvm") // 已经完成了数据特征工程的数据一种保存格式
      .load(inputPath)

    // 2、切分训练集和测试集：8:2
    val dfArr: Array[DataFrame] = libsvmDF.randomSplit(Array(0.8, 0.2))
    val trainDF: DataFrame = dfArr(0)
    val testDF: DataFrame = dfArr(1)

    // 3、选择合适的模型进行训练
    /**
     * 如何选择模型？
     * 有无label？
     * 有 -> 有监督学习
     * label是否连续?
     * 连续的 -> 回归算法：线性回归、岭回归
     * 离散的 -> 分类算法：KNN、决策树、逻辑回归、朴素贝叶斯分类
     * 无 -> 无监督学习
     */
    val lr: LogisticRegression = new LogisticRegression()
      .setMaxIter(10)
      .setRegParam(0.3)
      .setElasticNetParam(0.8)

    // Fit the model
    val lrModel: LogisticRegressionModel = lr.fit(trainDF)

    // 4、使用测试集进行模型的评估
    val resDF: DataFrame = lrModel.transform(testDF)

    resDF.cache()
    import spark.implicits._

    println(s"模型的准确率为：${resDF.where($"label" === $"prediction").count().toDouble / resDF.count()}")

    // 5、保存模型并使用
    lrModel
      .write
      .overwrite()
      .save(outputPath)


    resDF.unpersist()
  }

}
